package models

import (
	"context"

	"wx-uniapp-go/third_party/mysql"
	"xorm.io/xorm"
)

type ctxSessionKey = struct{}

var (
	getDataBelongEpId func(ctx context.Context) string
)

func RegisterGetEpId(fn func(ctx context.Context) string) {
	getDataBelongEpId = fn
}

func defaultGetSession(ctx context.Context) *xorm.Session {
	return mysql.DC(ctx)
}

// WithSession 兼容已经开过 session 的代码
func WithSession(ctx context.Context, ss *xorm.Session) context.Context {
	return setSession(ctx, ss)
}

func setSession(ctx context.Context, ss *xorm.Session) context.Context {
	return context.WithValue(ctx, ctxSessionKey{}, ss)
}

func getSession(ctx context.Context) *xorm.Session {
	ss, ok := ctx.Value(ctxSessionKey{}).(*xorm.Session)
	if !ok {
		return defaultGetSession(ctx)
	}
	return ss
}

func BeginTransaction(ctx context.Context) context.Context {
	ss := mysql.Default().NewSession().Context(ctx)
	ss.Begin()
	return setSession(ctx, ss)
}

func CloseTransaction(ctx context.Context, err error) error {
	ss, ok := ctx.Value(ctxSessionKey{}).(*xorm.Session)
	if !ok {
		return nil
	}
	defer ss.Close()
	if err != nil {
		return ss.Rollback()
	}
	return ss.Commit()
}
